from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session
from contextlib import contextmanager
from typing import Generator
from sqlalchemy.exc import SQLAlchemyError

from app.core.config import settings
from app.utils.log_utils import logger_util

logger = logger_util.get_logger(__name__)

# 创建数据库引擎
engine = create_engine(
    settings.mysql_url,
    pool_pre_ping=True,
    pool_recycle=3600,
    pool_size=5,
    max_overflow=10,
    echo=False
)

# 创建会话工厂
SessionLocal = sessionmaker(
    bind=engine,
    autocommit=False,
    autoflush=False,
    expire_on_commit=False
)


@contextmanager
def get_db() -> Generator[Session, None, None]:
    """数据库会话上下文管理器"""
    session = SessionLocal()
    try:
        yield session
        session.commit()
    except SQLAlchemyError as e:
        session.rollback()
        logger.error(f"Database error: {str(e)}")
        raise
    except Exception as e:
        session.rollback()
        logger.error(f"Non-database error: {str(e)}")
        raise
    finally:
        session.close()


# 用于测试数据库连接
def test_db_connection():
    try:
        with get_db() as db:
            # 使用PostgreSQL特定的查询语法
            db.execute("SELECT 1")
            logger.info("Database connection successful")
        return True
    except Exception as e:
        logger.error(f"Database connection failed: {str(e)}")
        return False
